#include <vector>

#include "caffe/layer.hpp"
#include "caffe/util/math_functions.hpp"
#include "caffe/util/interp.hpp"
#include "caffe/layers/interp_layer.hpp"

namespace caffe {

template <typename Dtype>
void InterpLayer<Dtype>::LayerSetUp(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top) {
  InterpParameter interp_param = this->layer_param_.interp_param();
  pad_beg_ = interp_param.pad_beg();
  pad_end_ = interp_param.pad_end();
  align_corners_ = interp_param.align_corners();
  half_pixel_center_ = interp_param.half_pixel_center();
  mode_ = this->layer_param_.interp_param().mode();
  CHECK_LE(pad_beg_, 0) << "Only supports non-pos padding (cropping) for now";
  CHECK_LE(pad_end_, 0) << "Only supports non-pos padding (cropping) for now";
}

template <typename Dtype>
void InterpLayer<Dtype>::Reshape(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top) {
  if (align_corners_) {
    num_ = bottom[0]->num();
    channels_ = bottom[0]->channels();
    height_in_ = bottom[0]->height();
    width_in_ = bottom[0]->width();
    height_in_eff_ = height_in_ + pad_beg_ + pad_end_;
    width_in_eff_ = width_in_ + pad_beg_ + pad_end_;
    InterpParameter interp_param = this->layer_param_.interp_param();
    if (interp_param.has_shrink_factor() &&
        !interp_param.has_zoom_factor()) {
      const int shrink_factor = interp_param.shrink_factor();
      CHECK_GE(shrink_factor, 1) << "Shrink factor must be positive";
      height_out_ = (height_in_eff_ - 0) / shrink_factor + 1;
      width_out_ = (width_in_eff_ - 0) / shrink_factor + 1;
    } else if (interp_param.has_zoom_factor() &&
              !interp_param.has_shrink_factor()) {
      const int zoom_factor = interp_param.zoom_factor();
      CHECK_GE(zoom_factor, 1) << "Zoom factor must be positive";
      height_out_ = height_in_eff_ + (height_in_eff_ - 0) * (zoom_factor - 1);
      width_out_ = width_in_eff_ + (width_in_eff_ - 0) * (zoom_factor - 1);
    } else if (interp_param.has_height() && interp_param.has_width()) {
      height_out_  = interp_param.height();
      width_out_  = interp_param.width();
    } else if (interp_param.has_shrink_factor() &&
              interp_param.has_zoom_factor()) {
      const int shrink_factor = interp_param.shrink_factor();
      const int zoom_factor = interp_param.zoom_factor();
      CHECK_GE(shrink_factor, 1) << "Shrink factor must be positive";
      CHECK_GE(zoom_factor, 1) << "Zoom factor must be positive";
      height_out_ = (height_in_eff_ - 0) / shrink_factor + 1;
      width_out_ = (width_in_eff_ - 0) / shrink_factor + 1;
      height_out_ = height_out_ + (height_out_ - 0) * (zoom_factor - 1);
      width_out_ = width_out_ + (width_out_ - 0) * (zoom_factor - 1);
    } else {
      LOG(FATAL);
    }
    CHECK_GT(height_in_eff_, 0) << "height should be positive";
    CHECK_GT(width_in_eff_, 0) << "width should be positive";
    CHECK_GT(height_out_, 0) << "height should be positive";
    CHECK_GT(width_out_, 0) << "width should be positive";
    top[0]->Reshape(num_, channels_, height_out_, width_out_);
  } else {
    num_ = bottom[0]->num();
    channels_ = bottom[0]->channels();
    height_in_ = bottom[0]->height();
    width_in_ = bottom[0]->width();
    InterpParameter interp_param =
      this->layer_param_.interp_param();
    if (interp_param.has_zoom_factor()) {
      const int zoom_factor = interp_param.zoom_factor();
      height_out_ = static_cast<int>(height_in_ * zoom_factor);
      width_out_ = static_cast<int>(width_in_ * zoom_factor);
    } else if (interp_param.has_height() && interp_param.has_width()) {
      height_out_ = interp_param.height();
      width_out_ = interp_param.width();
    } else {
      LOG(FATAL);
    }
    
    CHECK_GT(height_in_, 0) << "input height shoule be positive";
    CHECK_GT(width_in_, 0) << "input width shoule be positive";
    CHECK_GT(height_out_, 0) << "output height shoule be positive";
    CHECK_GT(width_out_, 0) << "output width shoule be positive";
    top[0]->Reshape(num_, channels_, height_out_, width_out_);
  }
}

template <typename Dtype>
void InterpLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top) {
  switch (mode_) {
    case InterpParameter_InterpOp_nearest:
    {
      std::cout<<"===========InterpParameter_InterpOp_nearest=============="<<std::endl;
      float height_scale = static_cast<float>(height_in_) / height_out_;
      float width_scale = static_cast<float>(width_in_) / width_out_;
      const Dtype *input_data = bottom[0]->cpu_data();
      Dtype *output_data = top[0]->mutable_cpu_data();
      for (int b = 0; b < num_; b++)
        {
          for (int c = 0; c < channels_; c++)
          {
            const Dtype *in_start = input_data + b * channels_ * height_in_ * width_in_ + c * height_in_ * width_in_;
            Dtype *out_start = output_data + b * channels_ * height_out_ * width_out_  + c * height_out_ * width_out_ ;
            for (int y = 0; y < height_out_; y++)
            {
              //int in_y = MIN((int)( floorf(y*height_scale)), height_in_ - 1);
              int in_y = MIN((int)(align_corners_ ? roundf(y*height_scale) : floorf(y*height_scale)), height_in_ - 1);
              for (int x = 0; x < width_out_; x++)
              {
                //int in_x = MIN((int)(floorf(x*width_scale)), width_in_ - 1);
               int in_x = MIN((int)(align_corners_ ? roundf(x*width_scale) : floorf(x*width_scale)), width_in_ - 1);
                Dtype interpolation = static_cast<Dtype>(in_start[in_y*width_in_ + in_x]);
                out_start[y*width_out_ + x] = interpolation ;
              }
            }
          }
        }
    }
    break;
    case InterpParameter_InterpOp_bilinear:
    {
      std::cout<<"===========InterpParameter_InterpOp_bilinear=============="<<std::endl;
      if (align_corners_)
       {
          std::cout<<"===========InterpParameter_AlignMethod_Corner=============="<<std::endl;
          caffe_cpu_interp2<Dtype,false>(num_ * channels_,
            bottom[0]->cpu_data(), - pad_beg_, - pad_beg_, height_in_eff_, width_in_eff_, height_in_, width_in_,
            top[0]->mutable_cpu_data(), 0, 0, height_out_, width_out_, height_out_, width_out_);
        } 
        else
         {
          std::cout<<"===========InterpParameter_AlignMethod_Center=============="<<std::endl;
          float height_scale = static_cast<float>(height_in_) / height_out_;
          float width_scale = static_cast<float>(width_in_) / width_out_;
          const Dtype *input_data = bottom[0]->cpu_data();
          Dtype *output_data = top[0]->mutable_cpu_data();
          float input_y, input_x;
          for (int b = 0; b < num_; ++b) {
            for (int y = 0; y < height_out_; ++y) {
              if(half_pixel_center_){
                 input_y = y * height_scale;//half_pixel_center = true
              }
              else{
                 input_y = ((float)(y)+0.5f) * height_scale-0.5;//half_pixel_center = false
              }

              int y0 = static_cast<int>(std::floor(input_y));
              int y1 = std::min(y0 + 1, height_in_ - 1);

              for (int x = 0; x < width_out_; ++x) {
                if(half_pixel_center_){
                   input_x = x * width_scale;
                }
                else{
                   input_x = ((float)(x)+0.5f) * width_scale-0.5;
                }
                
                int x0 = static_cast<int>(std::floor(input_x));
                int x1 = std::min(x0 + 1, width_in_ - 1);

                float sx = input_x - x0;
                float sy = input_y - y0;
                if(y0<0)
               		sy = 1;
                if(x0<0)
                	sx = 1;
                for (int c = 0; c < channels_; ++c) {
                  Dtype interpolation =
                    static_cast<Dtype>(input_data[b * channels_ * height_in_ *
                        width_in_ + c * height_in_ * width_in_ +
                        y0 * width_in_ + x0]
                          * (1 - sy) * (1 - sx) +
                        input_data[b * channels_ * height_in_ * width_in_ +
                        c * height_in_ * width_in_ + y1 * width_in_ + x0]
                          * sy * (1 - sx) +
                        input_data[b * channels_ * height_in_ * width_in_ +
                        c * height_in_ * width_in_ + y0 * width_in_ + x1]
                          * (1 - sy) * sx +
                        input_data[b * channels_ * height_in_ * width_in_ +
                        c * height_in_ * width_in_ + y1 * width_in_ + x1]
                          * sy * sx);
                  output_data[b * channels_ * height_out_ * width_out_ +
                    c * height_out_ * width_out_
                    + y * width_out_ + x] = interpolation;
                }


              }
            }
          }
        }
    }
    break;
      default:
    LOG(FATAL) << "Unknown Interp mode.";
  }
}

template <typename Dtype>
void InterpLayer<Dtype>::Backward_cpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom) {
  if (!propagate_down[0]) { return; }
    switch (mode_) {
    case InterpParameter_InterpOp_nearest:
      {
        NOT_IMPLEMENTED;
      }
    break;
    case InterpParameter_InterpOp_bilinear:
    {
      if (align_corners_) {
          caffe_set(bottom[0]->count(), Dtype(0), bottom[0]->mutable_cpu_diff());
          caffe_cpu_interp2_backward<Dtype,false>(num_ * channels_,
            bottom[0]->mutable_cpu_diff(), - pad_beg_, - pad_beg_, height_in_eff_, width_in_eff_, height_in_, width_in_,
            top[0]->cpu_diff(), 0, 0, height_out_, width_out_, height_out_, width_out_);
      } else {
          NOT_IMPLEMENTED;
      } 
    }
    break;
    default:
    LOG(FATAL) << "Unknown Interp mode.";
    }
}

#ifndef CPU_ONLY
template <typename Dtype>
void InterpLayer<Dtype>::Forward_gpu(const vector<Blob<Dtype>*>& bottom,
    const vector<Blob<Dtype>*>& top) {
    switch (mode_) {
    case InterpParameter_InterpOp_nearest:
      {
        NOT_IMPLEMENTED;
      }
    break;
    case InterpParameter_InterpOp_bilinear:
    {
    if (align_corners_) {
      caffe_gpu_interp2<Dtype,false>(num_ * channels_,
        bottom[0]->gpu_data(), - pad_beg_, - pad_beg_, height_in_eff_, width_in_eff_, height_in_, width_in_,
        top[0]->mutable_gpu_data(), 0, 0, height_out_, width_out_, height_out_, width_out_);
    } else {
      NOT_IMPLEMENTED;
    }
  }
    break;
    default:
    LOG(FATAL) << "Unknown Interp mode.";
    }
}

template <typename Dtype>
void InterpLayer<Dtype>::Backward_gpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom) {
  if (!propagate_down[0]) { return; }
  switch (mode_) {
    case InterpParameter_InterpOp_nearest:
      {
        NOT_IMPLEMENTED;
      }
    break;
    case InterpParameter_InterpOp_bilinear:
    {
      if (align_corners_) {
        caffe_gpu_set(bottom[0]->count(), Dtype(0), bottom[0]->mutable_gpu_diff());
        caffe_gpu_interp2_backward<Dtype,false>(num_ * channels_,
          bottom[0]->mutable_gpu_diff(), - pad_beg_, - pad_beg_, height_in_eff_, width_in_eff_, height_in_, width_in_,
          top[0]->gpu_diff(), 0, 0, height_out_, width_out_, height_out_, width_out_);
      } else {
        NOT_IMPLEMENTED;
      }
    }
    break;
    default:
    LOG(FATAL) << "Unknown Interp mode.";
    }
}
#endif

#ifdef CPU_ONLY
STUB_GPU(InterpLayer);
#endif

INSTANTIATE_CLASS(InterpLayer);
REGISTER_LAYER_CLASS(Interp);

} // namespace caffe